import os

root_dir = "./hymenoptera_data/hymenoptera_data/train/"

ants_dir = "ants_image"
bees_dir = "bees_image"
ants_img_path = os.listdir(os.path.join(root_dir, ants_dir))
bees_img_path = os.listdir(os.path.join(root_dir, bees_dir))

ants_label = ants_dir.split("_")[0]
bees_label = bees_dir.split('_')[0]

ant_out_dir = "ants_label/"
bees_out_dir = "bees_label/"

for i in ants_img_path:
    file_name = i.split('.jpg')[0]
    with open(os.path.join(root_dir, ant_out_dir, "{}.txt".format(file_name)), 'w') as f:
        f.write(ants_label)

for j in bees_img_path:
    file_name = j.split('.jpg')[0]
    with open(os.path.join(root_dir, bees_out_dir, "{}.txt".format(file_name)), 'w') as f:
        f.write(bees_label)